# encoding: utf-8
#!/usr/bin/python3
import numpy as np
import torch
from tqdm import tqdm
import torch.nn.functional as F
import cv2
from loss import multi_iou,mean_iou






def val_fun(model,data_loader):
    all_iou = np.zeros((1,10))
    model.eval()
    with torch.no_grad():
        for batch_image,batch_label in tqdm(data_loader):
            batch_pred = model(batch_image.float().cuda())
            batch_pred = torch.argmax(F.log_softmax(batch_pred, dim=1), dim=1)
            batch_pred = batch_pred.cpu().numpy()
            for batch_idx in range(batch_pred.shape[0]):
                iou = mean_iou(batch_pred[batch_idx],batch_label[batch_idx].cpu().numpy(),classes=10)
                # print(iou)
                all_iou = np.concatenate((all_iou,iou),axis=0)
    return np.nanmean(all_iou,axis=0)





if __name__ == '__main__':
    import torch
    from models.model import *
    from torch.utils.data import DataLoader
    from dataset import MyDataset

    Image_Path = '../suichang_round1_train_210120\suichang_round1_train_210120'
    Train_Json_Dir = 'label_file/train.json'
    Val_Json_Dir = 'label_file/val.json'


    model = model_Unet_resnet50
    model.load_state_dict(torch.load('../model_save/Unet_resnet50/51.pth'))
    val_dataSet = MyDataset(image_path=Image_Path, json_dir=Val_Json_Dir, data_aug=False)
    val_data_loader = DataLoader(dataset=val_dataSet, batch_size=2, shuffle=True,)
    iou = val_fun(model,val_data_loader)
    print(iou)
    print(np.nanmean(iou))






